#include "stdafx.h"

#include <gtest/gtest.h>
#include <ppl.h>

#include "../../../XProcess/JmpHooker.h"
#include "../../../XProcess/Module.h"
#include <future>
#include <vector>

using namespace XProcess;
using namespace Concurrency;
using namespace std;

volatile atomic<int> sum = 0;
void __stdcall incVar()
{
    for (int i = 0; i < 100 ; i++)
    {
        sum++;
    }
}

void __stdcall fixedIncVar(decltype(&incVar) real_incVar)
{
    for (int i = 0; i < 10; i++)
    {
        real_incVar();
    }
}
TEST(Hook, MultiThread)
{
    {
        vector<future<void>> fs;
        for (int i = 0; i < 100; i++)
        {
            fs.emplace_back(std::async(std::launch::async, &incVar));
        }

        bool finished = false;
        for (auto &f : fs)
        {
            f.wait();
        }
    }
    EXPECT_EQ(100 * 100, sum);


    sum = 0;
    // hook it
    auto hooker = JmpHooker<3, void()>::hook(&incVar, &fixedIncVar);
    {
        vector<future<void>> fs;
        for (int i = 0; i < 100; i++)
        {
            fs.emplace_back(std::async(std::launch::async, &incVar));
        }

        bool finished = false;
        for (auto &f : fs)
        {
            f.wait();
        }
    }
    EXPECT_EQ(10 * 100 * 100, sum);
    hooker.unHook();
}